import pandas as pd
import numpy as np
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from sklearn.impute import KNNImputer                     # >>> 新增
from statsmodels.tsa.holtwinters import SimpleExpSmoothing  # >>> 新增
import os

# 缺失值处理记录
fill_method_records = []

def fill_missing_values(df, country, indicator, method):
    """根据不同的方法填充缺失值"""
    if method == '1':  # 均值填充
        fill_value = df[indicator].mean()
        df[indicator].fillna(fill_value, inplace=True)

    elif method == '2':  # 前向填充
        df[indicator].fillna(method='ffill', inplace=True)

    elif method == '3':  # 后向填充
        df[indicator].fillna(method='bfill', inplace=True)

    elif method == '4':  # 前后向填充
        if pd.isna(df[indicator].iloc[0]):
            df[indicator].fillna(method='bfill', inplace=True)
        df[indicator].fillna(method='ffill', inplace=True)

    elif method == '5':  # 线性插值
        if pd.isna(df[indicator].iloc[0]):
            df[indicator].fillna(method='bfill', inplace=True)
        df[indicator].interpolate(method='linear', inplace=True)

    elif method == '6':  # 多项式插值（三阶）
        if df[indicator].isna().iloc[0]:
            df[indicator].fillna(method='bfill', inplace=True)
            if df[indicator].isna().iloc[0]:
                df[indicator].fillna(method='ffill', inplace=True)

        df_non_nan = df.dropna(subset=[indicator])
        X = df_non_nan[['Year']]
        y = df_non_nan[indicator]
        poly = PolynomialFeatures(degree=3)
        X_poly = poly.fit_transform(X)
        poly_reg = LinearRegression()
        poly_reg.fit(X_poly, y)

        for idx in df.index:
            if pd.isna(df.loc[idx, indicator]):
                X_pred = poly.transform([[df.loc[idx, 'Year']]])
                pred = poly_reg.predict(X_pred)
                df.loc[idx, indicator] = max(pred[0], 0)

    elif method == '7':  # KNN 插值             # >>> 新增
        knn_df = df[['Year', indicator]].copy()  # >>> 新增
        imputer = KNNImputer(n_neighbors=3)      # >>> 新增
        imputed = imputer.fit_transform(knn_df)  # >>> 新增
        df[indicator] = imputed[:, 1]            # >>> 新增

    elif method == '8':  # 指数平滑             # >>> 新增
        series = df.sort_values('Year')[indicator]   # >>> 新增
        if series.notna().sum() >= 3:                # >>> 新增
            model = SimpleExpSmoothing(series.dropna(), initialization_method="heuristic").fit()  # >>> 新增
            filled_series = series.copy()            # >>> 新增
            filled_series[series.isna()] = model.predict(series.index[series.isna()])  # >>> 新增
            df[indicator] = filled_series            # >>> 新增
        else:                                         # >>> 新增
            df[indicator].fillna(method='ffill', inplace=True)  # >>> 新增
            df[indicator].fillna(method='bfill', inplace=True)  # >>> 新增

    return df

def process_missing_values(df, index_list, country_list, output_root):
    """对指定的指标和国家进行缺失值处理"""
    global fill_method_records
    for indicator in index_list:
        record_data = []
        for i, country in enumerate(country_list, start=1):
            country_df = df[df['国名Ch'] == country][['国名Ch', indicator, 'Year']].copy()
            stats = country_df[indicator].describe()
            missing_percent = country_df[indicator].isna().mean() * 100

            print(f"\n当前指标数据（{indicator} - {country}）：")
            print(country_df)
            print(f"当前指标描述统计结果：\n{stats}")
            print(f"当前指标缺失值百分比：{missing_percent:.2f}%")

            user_input = input(f"是否需要对指标 {indicator}、国家 {country} 进行缺失值处理？(Y/N): ").strip().upper()
            if user_input == 'N':
                continue

            while True:
                method = input(
                    "请选择缺失值处理方法：\n"
                    " 1. 均值填充\n"
                    " 2. 前向填充\n"
                    " 3. 后向填充\n"
                    " 4. 前后向填充\n"
                    " 5. 线性插值\n"
                    " 6. 多项式插值（三阶）\n"
                    " 7. KNN 插值\n"         # >>> 新增
                    " 8. 指数平滑\n"         # >>> 新增
                    "请输入: "
                ).strip()
                if method in [str(i) for i in range(1, 9)]:   # >>> 修改：从 1-6 改为 1-8
                    break
                else:
                    print("无效输入，请重新输入 1-8 之间的数字。")  # >>> 修改提示内容

            method_dict = {
                '1': '均值填充',
                '2': '前向填充',
                '3': '后向填充',
                '4': '前后向填充',
                '5': '线性插值',
                '6': '多项式插值',
                '7': 'KNN 插值',         # >>> 新增
                '8': '指数平滑'          # >>> 新增
            }
            record_data.append([i, country, method_dict[method]])

            df.loc[df['国名Ch'] == country, indicator] = fill_missing_values(
                df.loc[df['国名Ch'] == country], country, indicator, method)[indicator]

            print(f"处理后数据（{indicator} - {country}）：")
            print(df[df['国名Ch'] == country][['国名Ch', indicator, 'Year']])

        fill_method_records.extend(record_data)

    # 保存缺失值处理方法记录
    record_df = pd.DataFrame(fill_method_records, columns=['序号', '国家', '处理方法'])
    os.makedirs(output_root, exist_ok=True)
    record_df.to_excel(os.path.join(output_root, '填充方法记录.xlsx'), index=False)
    print(f"\n缺失值处理方法记录已保存至 {output_root}/填充方法记录.xlsx")


#使用注释 # >>> 新增 标出新增或修改的部分